"""
Prompt Manager 模块

用于管理和加载各种操作的 prompt 模板，支持动态替换占位符。
这使得用户可以通过修改 prompt 文件来控制大模型的操作方式，
而无需修改代码。
"""

import os
import logging
from string import Template
from typing import Dict, Any, List, Optional
import json
import re

logger = logging.getLogger(__name__)

class PromptManager:
    """Prompt 管理器，用于加载和处理各种操作的 prompt 模板"""
    
    def __init__(self, prompt_dir: str = "prompts"):
        """
        初始化 Prompt 管理器
        
        参数:
            prompt_dir: prompt 模板文件所在的目录，默认为 'prompts'
        """
        self.prompt_dir = prompt_dir
        self.prompts = {}
        self._load_prompts()
    
    def _load_prompts(self) -> None:
        """加载所有 prompt 模板文件"""
        try:
            if not os.path.exists(self.prompt_dir):
                os.makedirs(self.prompt_dir)
                logger.warning(f"Prompt 目录 '{self.prompt_dir}' 不存在，已创建")
                
            prompt_files = [f for f in os.listdir(self.prompt_dir) if f.endswith('.txt')]
            
            if not prompt_files:
                logger.warning(f"在 '{self.prompt_dir}' 目录中没有找到任何 prompt 模板文件")
                return
                
            for filename in prompt_files:
                prompt_name = os.path.splitext(filename)[0]
                file_path = os.path.join(self.prompt_dir, filename)
                
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        self.prompts[prompt_name] = f.read()
                    logger.info(f"已加载 prompt 模板: {prompt_name}")
                except Exception as e:
                    logger.error(f"加载 prompt 模板 '{prompt_name}' 时出错: {str(e)}")
        except Exception as e:
            logger.error(f"加载 prompt 模板文件时出错: {str(e)}")
    
    def get_prompt(self, prompt_name: str, **kwargs) -> Optional[str]:
        """
        获取指定名称的 prompt 模板，并填充占位符
        
        参数:
            prompt_name: prompt 模板名称（不含扩展名）
            **kwargs: 用于替换模板中占位符的键值对
            
        返回:
            填充好占位符的 prompt 字符串，如果模板不存在则返回 None
        """
        if prompt_name not in self.prompts:
            logger.warning(f"Prompt 模板 '{prompt_name}' 不存在")
            return None
            
        try:
            # 记录传入的参数
            logger.info(f"准备处理 prompt 模板 '{prompt_name}'")
            logger.info(f"传入的参数: {json.dumps(kwargs, ensure_ascii=False, indent=2)}")
            
            template_content = self.prompts[prompt_name]
            
            # 检查模板中的所有占位符
            placeholders = set(re.findall(r'\{(\w+)\}', template_content))
            missing_params = placeholders - set(kwargs.keys())
            
            if missing_params:
                logger.error(f"缺少必要的参数: {missing_params}")
                return None
                
            # 使用 format 方法进行替换
            try:
                result = template_content.format(**kwargs)
            except KeyError as e:
                logger.error(f"参数替换失败，缺少参数: {str(e)}")
                return None
            except ValueError as e:
                logger.error(f"参数替换失败，格式错误: {str(e)}")
                return None
            
            # 记录替换后的结果预览
            preview = result[:200] + "..." if len(result) > 200 else result
            logger.info(f"生成的 prompt 预览: {preview}")
            
            return result
        except Exception as e:
            logger.error(f"处理 prompt 模板 '{prompt_name}' 时出错: {str(e)}")
            return None
    
    def list_available_prompts(self) -> List[str]:
        """
        获取所有可用的 prompt 模板名称列表
        
        返回:
            prompt 模板名称列表
        """
        return list(self.prompts.keys())
    
    def reload_prompts(self) -> None:
        """重新加载所有 prompt 模板文件"""
        self.prompts = {}
        self._load_prompts()


# 使用示例
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # 初始化 Prompt 管理器
    manager = PromptManager()
    
    # 列出所有可用的 prompt 模板
    available_prompts = manager.list_available_prompts()
    print(f"可用的 prompt 模板: {available_prompts}")
    
    # 获取翻译 prompt 示例
    translate_prompt = manager.get_prompt(
        "translate", 
        target_language="英文",
        source_text="这是一个示例文本，用于测试翻译功能。"
    )
    print("\n翻译 Prompt 示例:")
    print(translate_prompt)
    
    # 获取总结 prompt 示例
    summarize_prompt = manager.get_prompt(
        "summarize", 
        target_language="中文",
        summary_length=200,
        source_text="这是一个较长的示例文本，用于测试总结功能。可以包含多个段落和详细信息。"
    )
    print("\n总结 Prompt 示例:")
    print(summarize_prompt) 